#!/usr/bin/env python3
import os
import sys
import time
from collections import Counter, deque
from typing import Deque, Dict, Optional, Tuple

import rospy
from std_msgs.msg import String


def _ensure_paths() -> None:
    """Ensure we can import detection pipeline and GStreamer wrapper."""
    # scripts -> yolo_bridge -> src -> nav_ws -> repo_root
    root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
    yolo_detect_dir = os.path.join(root, 'yolo-jetson', 'detect')
    yolo_rs2_dir = os.path.join(root, 'yolo-jetson', 'rs2')
    if yolo_detect_dir not in sys.path:
        sys.path.append(yolo_detect_dir)
    if yolo_rs2_dir not in sys.path:
        sys.path.append(yolo_rs2_dir)


_ensure_paths()

# Now safe to import
from video_stream_optimized import VideoStreamPipelineOptimized  # type: ignore
from GStreamerWrapper import GStreamerWrapper  # type: ignore


class YOLOBridgeNode:
    def __init__(self) -> None:
        self._result_pub = rospy.Publisher('/yolo_result', String, queue_size=10)
        self._control_sub = rospy.Subscriber('/yolo_detect', String, self._control_cb, queue_size=10)

        # Params
        # Allow headless matplotlib usage if imported by pipeline
        os.environ.setdefault('MPLBACKEND', 'Agg')

        self.model_path: str = rospy.get_param('~model_path', r".\\weights\\all_yolo_all_epochs_0811\\best.pt")
        self.width: int = int(rospy.get_param('~width', 1920))
        self.height: int = int(rospy.get_param('~height', 1080))
        self.fps_log_interval: float = float(rospy.get_param('~fps_log_interval', 2.0))
        self.use_fresh_pipeline_per_frame: bool = bool(rospy.get_param('~use_fresh_pipeline_per_frame', True))
        self.required_window: int = int(rospy.get_param('~required_window', 5))

        # State
        self._running: bool = False
        self._gs: Optional[GStreamerWrapper] = None
        self._pipeline: Optional[VideoStreamPipelineOptimized] = None
        self._candidate_window: Deque[Tuple[str, str, str]] = deque(maxlen=self.required_window)
        self._frame_count: int = 0
        self._timer_last: float = time.time()
        self._timer_frames: int = 0

        rospy.loginfo('yolo_rtsp_node ready. Subscribe /yolo_detect and publish /yolo_result')

    def _init_gs_and_pipeline(self) -> None:
        if self._gs is None:
            self._gs = GStreamerWrapper(width=self.width, height=self.height)
            rospy.loginfo('GStreamer RTSP initialized')
        if self._pipeline is None:
            model_path = self._resolve_model_path(self.model_path)
            self._pipeline = VideoStreamPipelineOptimized(
                model_path,
                use_fresh_pipeline_per_frame=self.use_fresh_pipeline_per_frame,
            )
            rospy.loginfo('YOLO optimized pipeline initialized')

    def _shutdown_gs_and_pipeline(self) -> None:
        if self._gs is not None:
            try:
                self._gs.StopThread()
            except Exception:
                pass
            self._gs = None
        if self._pipeline is not None:
            try:
                del self._pipeline
            except Exception:
                pass
            self._pipeline = None
        self._candidate_window.clear()
        self._frame_count = 0
        self._timer_last = time.time()
        self._timer_frames = 0

    def _control_cb(self, msg: String) -> None:
        cmd = (msg.data or '').strip().lower()
        if cmd == 'start' and not self._running:
            rospy.loginfo('Received start. Begin detection.')
            self._running = True
            self._candidate_window.clear()
            self._frame_count = 0
        elif cmd == 'stop' and self._running:
            rospy.loginfo('Received stop. Halt detection.')
            self._running = False
            self._shutdown_gs_and_pipeline()

    @staticmethod
    def _resolve_model_path(cfg_path: str) -> str:
        if cfg_path and os.path.exists(cfg_path):
            return cfg_path
        # search common candidates relative to repo root
        root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
        candidates = [
            os.path.join(root, 'yolo-jetson', 'detect', 'weights', 'all_yolo_all_epochs_0811', 'best.pt'),
            os.path.join(root, 'yolo-jetson', 'detect', 'weights', 'all_yolo_all_epochs_0808', 'best.pt'),
            os.path.join(root, 'yolo-jetson', 'detect', 'weights', 'all_yolo_all_epochs_0805', 'best.pt'),
            os.path.join(root, 'yolo-jetson', 'detect', 'weights', 'all_yolo_all_epochs_0804', 'best.pt'),
        ]
        for p in candidates:
            if os.path.exists(p):
                return p
        # fallback to provided string
        return cfg_path

    @staticmethod
    def _format_triplet(color: str, number: str, dashboard: str) -> str:
        return f"{color}-{number}-{dashboard}"

    @staticmethod
    def _extract_triplet(result_dict: Dict) -> Optional[Tuple[str, str, str]]:
        color = str(result_dict.get('cone_color', 'None'))
        number = str(result_dict.get('number', 'None'))
        dashboard = str(result_dict.get('dashboard_state', 'None'))
        if color != 'None' and number != 'None' and dashboard != 'None':
            return (color, number, dashboard)
        return None

    @staticmethod
    def _vote_majority(triplets: Deque[Tuple[str, str, str]]) -> Tuple[str, str, str]:
        # majority by whole triplet
        counts = Counter(triplets)
        return counts.most_common(1)[0][0]

    def spin(self) -> None:
        rate_hz = float(rospy.get_param('~loop_rate', 60.0))
        rate = rospy.Rate(rate_hz)
        try:
            while not rospy.is_shutdown():
                if not self._running:
                    rate.sleep()
                    continue

                try:
                    if self._gs is None or self._pipeline is None:
                        self._init_gs_and_pipeline()

                    frame = self._gs.GetFrame() if self._gs is not None else None
                    if frame is None:
                        time.sleep(0.005)
                        continue

                    results = self._pipeline.process_single_frame_optimized(frame, self._frame_count)

                    triplet = self._extract_triplet(results)
                    if triplet is not None:
                        self._candidate_window.append(triplet)
                        rospy.loginfo(f"candidate {len(self._candidate_window)}/{self.required_window}: {self._format_triplet(*triplet)}")

                    # FPS logging only
                    self._timer_frames += 1
                    now = time.time()
                    if now - self._timer_last >= self.fps_log_interval:
                        fps = self._timer_frames / max(1e-6, (now - self._timer_last))
                        rospy.loginfo(f"processing FPS: {fps:.1f}")
                        self._timer_last = now
                        self._timer_frames = 0

                    # Settlement when window full
                    if len(self._candidate_window) >= self.required_window:
                        final_triplet = self._vote_majority(self._candidate_window)
                        final_str = self._format_triplet(*final_triplet)
                        rospy.loginfo(f"final: {final_str}")
                        self._result_pub.publish(String(data=final_str))
                        # stop and wait for next start
                        self._running = False
                        self._shutdown_gs_and_pipeline()

                    self._frame_count += 1

                except Exception as e:
                    rospy.logerr(f"detection error: {e}")
                    time.sleep(0.01)

                rate.sleep()
        finally:
            self._shutdown_gs_and_pipeline()


def main() -> None:
    rospy.init_node('yolo_rtsp_node', anonymous=False)
    node = YOLOBridgeNode()
    node.spin()


if __name__ == '__main__':
    main()


